use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::mem;

/// A trait for types that can be used in the set.
pub trait SetValue: Hash + Eq {
    /// A function to free the value, if necessary.
    fn free(&self) {}
}

/// A trait for hashing values.
pub trait SetHashFunc {
    fn hash(&self, value: &dyn SetValue) -> u64;
}

/// A trait for equality comparison.
pub trait SetEqualFunc {
    fn equals(&self, a: &dyn SetValue, b: &dyn SetValue) -> bool;
}

/// A trait for freeing values.
pub trait SetFreeFunc {
    fn free(&self, value: Box<dyn SetValue>);
}

/// A default hasher implementation.
pub struct DefaultSetHasher;

impl SetHashFunc for DefaultSetHasher {
    fn hash(&self, value: &dyn SetValue) -> u64 {
        let mut hasher = DefaultHasher::new();
        value.hash(&mut hasher);
        hasher.finish()
    }
}

/// A default equality function implementation.
pub struct DefaultSetEqualFunc;

impl SetEqualFunc for DefaultSetEqualFunc {
    fn equals(&self, a: &dyn SetValue, b: &dyn SetValue) -> bool {
        a == b
    }
}

/// A default free function implementation.
pub struct DefaultSetFreeFunc;

impl SetFreeFunc for DefaultSetFreeFunc {
    fn free(&self, _value: Box<dyn SetValue>) {}
}

/// A set entry.
struct SetEntry {
    data: Box<dyn SetValue>,
    next: Option<Box<SetEntry>>,
}

/// A hash set implementation.
pub struct Set {
    table: Vec<Option<Box<SetEntry>>>,
    entries: usize,
    table_size: usize,
    prime_index: usize,
    hash_func: Box<dyn SetHashFunc>,
    equal_func: Box<dyn SetEqualFunc>,
    free_func: Box<dyn SetFreeFunc>,
}

impl Set {
    /// Creates a new set.
    pub fn new(
        hash_func: Box<dyn SetHashFunc>,
        equal_func: Box<dyn SetEqualFunc>,
        free_func: Box<dyn SetFreeFunc>,
    ) -> Self {
        let mut set = Set {
            table: Vec::new(),
            entries: 0,
            table_size: 0,
            prime_index: 0,
            hash_func,
            equal_func,
            free_func,
        };
        set.allocate_table();
        set
    }

    /// Frees the set and all its entries.
    pub fn free(&mut self) {
        for entry in self.table.iter_mut() {
            while let Some(mut e) = entry.take() {
                self.free_func.free(e.data);
                *entry = e.next.take();
            }
        }
    }

    /// Registers a free function for the set.
    pub fn register_free_function(&mut self, free_func: Box<dyn SetFreeFunc>) {
        self.free_func = free_func;
    }

    /// Inserts a value into the set.
    pub fn insert(&mut self, data: Box<dyn SetValue>) -> bool {
        if (self.entries * 3) / self.table_size > 0 {
            self.enlarge();
        }

        let index = (self.hash_func.hash(&*data) % self.table_size as u64) as usize;

        if let Some(entry) = self.table[index].as_mut() {
            let mut rover = entry;
            while let Some(ref mut next) = rover.next {
                if self.equal_func.equals(&*data, &*next.data) {
                    return false;
                }
                rover = next;
            }
            rover.next = Some(Box::new(SetEntry { data, next: None }));
        } else {
            self.table[index] = Some(Box::new(SetEntry { data, next: None }));
        }

        self.entries += 1;
        true
    }

    /// Removes a value from the set.
    pub fn remove(&mut self, data: &dyn SetValue) -> bool {
        let index = (self.hash_func.hash(data) % self.table_size as u64) as usize;

        let mut prev = None;
        let mut rover = self.table[index].as_mut();

        while let Some(entry) = rover {
            if self.equal_func.equals(data, &*entry.data) {
                if let Some(prev) = prev {
                    prev.next = entry.next.take();
                } else {
                    self.table[index] = entry.next.take();
                }
                self.entries -= 1;
                self.free_func.free(entry.data);
                return true;
            }
            prev = Some(entry);
            rover = entry.next.as_mut();
        }

        false
    }

    /// Queries if a value is in the set.
    pub fn query(&self, data: &dyn SetValue) -> bool {
        let index = (self.hash_func.hash(data) % self.table_size as u64) as usize;

        let mut rover = self.table[index].as_ref();
        while let Some(entry) = rover {
            if self.equal_func.equals(data, &*entry.data) {
                return true;
            }
            rover = entry.next.as_ref();
        }

        false
    }

    /// Returns the number of entries in the set.
    pub fn num_entries(&self) -> usize {
        self.entries
    }

    /// Converts the set to an array.
    pub fn to_array(&self) -> Vec<Box<dyn SetValue>> {
        let mut array = Vec::with_capacity(self.entries);

        for entry in self.table.iter() {
            let mut rover = entry.as_ref();
            while let Some(e) = rover {
                array.push(e.data.clone());
                rover = e.next.as_ref();
            }
        }

        array
    }

    /// Enlarges the set's table.
    fn enlarge(&mut self) -> bool {
        let old_table = mem::replace(&mut self.table, Vec::new());
        let old_table_size = self.table_size;
        self.prime_index += 1;
        self.allocate_table();

        for mut entry in old_table {
            while let Some(mut e) = entry.take() {
                let index = (self.hash_func.hash(&*e.data) % self.table_size as u64) as usize;
                e.next = self.table[index].take();
                self.table[index] = Some(e);
                entry = e.next.take();
            }
        }

        true
    }

    /// Allocates the set's table.
    fn allocate_table(&mut self) {
        self.table_size = if self.prime_index < SET_PRIMES.len() {
            SET_PRIMES[self.prime_index]
        } else {
            self.entries * 10
        };
        self.table = vec![None; self.table_size];
    }
}

/// A set of good hash table prime numbers.
const SET_PRIMES: [usize; 24] = [
    193, 389, 769, 1543, 3079, 6151, 12289, 24593, 49157, 98317, 196613, 393241, 786433, 1572869,
    3145739, 6291469, 12582917, 25165843, 50331653, 100663319, 201326611, 402653189, 805306457,
    1610612741,
];

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_set_insert_remove() {
        let mut set = Set::new(
            Box::new(DefaultSetHasher),
            Box::new(DefaultSetEqualFunc),
            Box::new(DefaultSetFreeFunc),
        );

        let value1 = Box::new(1) as Box<dyn SetValue>;
        let value2 = Box::new(2) as Box<dyn SetValue>;

        assert!(set.insert(value1.clone()));
        assert!(set.insert(value2.clone()));
        assert_eq!(set.num_entries(), 2);

        assert!(set.query(&*value1));
        assert!(set.query(&*value2));

        assert!(set.remove(&*value1));
        assert!(!set.query(&*value1));
        assert_eq!(set.num_entries(), 1);

        assert!(set.remove(&*value2));
        assert!(!set.query(&*value2));
        assert_eq!(set.num_entries(), 0);
    }
}
